"""Predictor"""

import numpy as np
import matplotlib.pyplot as plt

class Predictor(object):
    """
    Predictor using estimated coefficients

    Attributes
    ----------
    steps : int
        Number of prediction steps
    true_data : np.array
        True data to compare the predicted with
    predicted_data : np.array
        Predicted data
    order : int
        Order of rewgr. model


    Methods
    -------
    __init__(steps, true_data=None)
        Constructor
    predict(est_theta, vec_data, steps=None)
        Prediction
    plot(true_data=None, predicted_data=None)
        Plot predicted vs. true data

    Note 
    ----
    vec_data must have '1' at the end if multiple steps-ahead prediction
    with abs. term is evaluated!!! The method is not designed for
    predictions w/o it!!!
    """

    def __init__(self, steps=1, true_data=None):
        self.steps = steps
        self.true_data = true_data
        self.predicted_data = []
        self.order = 1
        self._statistics_ready = False
        self.stats = {}

    def predict(self, est_theta, vec_data, steps=None):
        """
        Predictor function

        Parameters
        ----------
        est_theta : np.array
            Estimates of regr. coefficients
        vec_data : np.array
            Regression vector
        steps : int
            Number of pred. steps. Default is self.steps

        Returns
        -------
        pred : float
            Predicted data
        """
        self.order = len(vec_data)

        if steps != None:
        	self.steps = steps
        
        for i in xrange(self.steps):
        	pred = np.dot(est_theta, vec_data)
        	vec_data = np.hstack((np.array([pred]), 
        	                      vec_data[0:-2],
        	                      np.ones(1)))
       	self.predicted_data.append(pred)
        return pred

    def plot(self, true_data=None, predicted_data=None):
        """
        Plots true vs. predicted data and prints statistics

        Parameters
        ----------
        true_data : np.array
            True data. If None, self.predicted_data are used.
        predicted_data : np.array
            Predicted data. If None, self.predicted_data are used.
        """
        if not self._statistics_ready == True:
            self._prepare(true_data, predicted_data)

        fig = plt.figure(2)
        ax = fig.add_subplot(111)
        ax.grid(True)
        ax.hold(True)
        ax.plot(self.true_data, '-', color='grey')
        ax.plot(self.predicted_data, '-', color='black', lw='2')
        ax.set_title("Predictions")
        #plt.show()


    def print_statistics(self, true_data=None, predicted_data=None):
        """
        Prints statistics of prediction.
        """
        self.get_statistics()

        # statistics
        print "Mean err.:   ", self.stats["mean_err"]
        print "Var err.:    ", self.stats["var_err"]
        print "Std err.:    ", self.stats["std_err"]
        print "Median err.: ", self.stats["median_err"]
        print "Max err.:    ", self.stats["max_err"]
        print "Max -err.:   ", self.stats["max_neg_err"]
        print "MSE:         ", self.stats["mse"]
        print "RMSE:        ", self.stats["rmse"]


    def get_statistics(self, true_data=None, predicted_data=None):
        """
        Returns selected statistics.
        """
        if not self._statistics_ready == True:
            self._prepare(true_data, predicted_data)

        errors = self.true_data - self.predicted_data
        mse = 1./(len(errors) + 1.) * np.sum(errors ** 2)

        self.stats = {
            "mean_err": np.mean(errors),
            "var_err": np.var(errors),
            "std_err": np.std(errors),
            "median_err": np.median(errors),
            "max_err": np.max(errors),
            "max_neg_err": np.min(errors),
            "mse": mse,
            "rmse": np.sqrt(mse)
        }
        return self.stats
        



    def _prepare(self, true_data, predicted_data):
        """
        Prepare statistics
        """
        self._statistics_ready = True
        if predicted_data != None:
        	self.predicted_data = predicted_data 
        if true_data != None:
        	self.true_data = true_data
        # Trim the data to fit them
        self.predicted_data = np.hstack((np.zeros(self.order + self.steps - 1),
                                         np.array(self.predicted_data[:-(self.steps + self.order - 1)])))

        





#=================================================================
if __name__ == '__main__':
	est_theta = np.array([0.3, 0.4, 0.8])
	vec_data = np.array([1., 2., 1.])

	predictor = Predictor()
	print predictor.predict(est_theta, vec_data, 3)
